package com.qiongqi.datadir.utils.security;


import paillierp.Paillier;
import paillierp.PaillierThreshold;
import paillierp.key.KeyGen;
import paillierp.key.PaillierPrivateThresholdKey;
import paillierp.zkp.DecryptionZKP;

import javax.crypto.Cipher;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.util.Random;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static org.apache.commons.lang3.ArrayUtils.isNotEmpty;
import static org.apache.commons.lang3.StringUtils.trim;

public final class HEUtils {
    public static PaillierPrivateThresholdKey[] createThresholdPaillierKey(int partsNum){
        Random rnd = new SecureRandom();
        PaillierPrivateThresholdKey[] keys =
                KeyGen.PaillierThresholdKey(128, partsNum, partsNum, rnd.nextLong());
        return keys;
    }

    public static String paillierEncrypt(Paillier paillier, String fieldValue){
        BigInteger val = BigInteger.valueOf(Math.round(Double.parseDouble(trim(fieldValue)) * 1000000));
        return paillier.encrypt(val).toString();
    }

    public static String paillierDecrypt(PaillierPrivateThresholdKey[] prvKeys, String fieldValue){
        String decryptFieldValue = null;
        if(isNotEmpty(prvKeys)){
            PaillierThreshold[] ptArray = Stream.of(prvKeys)
                    .map(PaillierThreshold::new) // 等价于 key -> new PaillierThreshold(key)
                    .toArray(PaillierThreshold[]::new);
            DecryptionZKP[] dzs = IntStream.range(0, ptArray.length)
                    .mapToObj(idx -> ptArray[idx].decryptProof(new BigInteger(fieldValue))) // 生成解密证明
                    .toArray(DecryptionZKP[]::new);
            BigDecimal divider = new BigDecimal(BigInteger.valueOf(1000000));
            BigDecimal resultBigDecimal = new BigDecimal(ptArray[0].combineShares(dzs)).divide(divider, 2, java.math.RoundingMode.HALF_UP);
            // 5. 将结果转换为 double
            decryptFieldValue = resultBigDecimal.toPlainString();
        }
        return decryptFieldValue;
    }

    public static KeyPair createThresholdRsaKey() throws Exception {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
        keyGen.initialize(1024);
        return keyGen.generateKeyPair();
    }

    public static byte[] rsaEncrypt(PublicKey publicKey, byte[] data) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        return cipher.doFinal(data);
    }
    public static String  rsaDecrypt(PrivateKey privateKey, byte[] encryptedData) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        // 2. 执行解密操作[1,4](@ref)
        byte[] decryptedBytes = cipher.doFinal(encryptedData);

        // 3. 转换为UTF-8字符串[2,6](@ref)
        return new String(decryptedBytes, StandardCharsets.UTF_8);
    }

    public static void main(String[] args){
        PaillierPrivateThresholdKey[] keys = createThresholdPaillierKey(2);

        PaillierThreshold p1 = new PaillierThreshold(keys[0]);
        PaillierThreshold p2 = new PaillierThreshold(keys[1]);

        System.out.println(" Alice is given the public key .");
        Paillier alice = new Paillier(keys[0].extractPublicKey());
        // Alice encrypts a message
        BigInteger msg = BigInteger.valueOf(135819283);
        BigInteger emsg = alice.encrypt(msg);
        System.out.println(" Alice encrypts the message " + msg + " and sends " +
                emsg + " to everyone .");
        BigInteger msg2 = BigInteger.valueOf(135819283);
        BigInteger emsg2 = alice.encrypt(msg);
        emsg = emsg.add(emsg2);

        System.out.println(" p1 receives the message and tries to decrypt all alone :");
        BigInteger p1decrypt = p1.decryptOnly(emsg);
        if (p1decrypt.equals(msg)) {
            System.out.println(" p1 succeeds decrypting the message all alone .");
        } else {
            System.out.println(" p1 fails decrypting the message all alone.:(");
        }

        System.out.println(" p2 receive the message and create a partial decryptions .");
        DecryptionZKP p1share = p1.decryptProof(emsg);
        DecryptionZKP p2share = p2.decryptProof(emsg);

        try {
            BigInteger p2decrypt = p2.combineShares(p1share, p2share);
            if (p2decrypt.equals(msg)) {
                System.out.println(" p2 succeeds decrypting the message with p1 .");
            } else {
                System.out.println(" p2 fails decrypting the message with p1.:(");
            }
        } catch (IllegalArgumentException e) {
            System.out.println(" p2 fails decrypting and throws an error ");
        }

    }
}
